import sys
from rknn.api import RKNN

DATASET_PATH = '../../../datasets/COCO/coco_subset_20.txt'
DEFAULT_RKNN_PATH = './yolov8.rknn'
DEFAULT_QUANT = False

def parse_arg():
    if len(sys.argv) < 3:
        print("Usage: python3 {} onnx_model_path [platform] [dtype(optional)] [output_rknn_path(optional)]".format(sys.argv[0]));
        print("       platform choose from [rk3562, rk3566, rk3568, rk3576, rk3588, rv1126b, rv1109, rv1126, rk1808]")
        print("       dtype choose from [i8, fp] for [rk3562, rk3566, rk3568, rk3576, rk3588, rv1126b]")
        print("       dtype choose from [u8, fp] for [rv1109, rv1126, rk1808]")
        exit(1)

    model_path = sys.argv[1]
    platform = sys.argv[2]

    do_quant = DEFAULT_QUANT
    if len(sys.argv) > 3:
        model_type = sys.argv[3]
        if model_type not in ['i8', 'u8', 'fp']:
            print("ERROR: Invalid model type: {}".format(model_type))
            exit(1)
        elif model_type in ['i8', 'u8']:
            do_quant = True
        else:
            do_quant = False

    if len(sys.argv) > 4:
        output_path = sys.argv[4]
    else:
        output_path = DEFAULT_RKNN_PATH

    return model_path, platform, do_quant, output_path

if __name__ == '__main__':
    model_path, platform, do_quant, output_path = parse_arg()

    # Create RKNN object
    rknn = RKNN(verbose=True)

    # Pre-process config - 针对YOLOv8的特殊配置
    print('--> Config model')
    rknn.config(
        mean_values=[[0, 0, 0]], 
        std_values=[[255, 255, 255]], 
        target_platform=platform,
        batch_size=1,
        input_format='NHWC',
        # 禁用对YOLO模型的特殊处理
        model_pruning=False,
        # 使用更兼容的配置
        optimization_level=1,  # 降低优化级别以避免兼容性问题
        # 禁用可能引起问题的特性
        remove_weight=False,
        custom_string='yolov8',  # 明确指定模型类型
    )
    print('done')

    # Load model
    print('--> Loading model')
    ret = rknn.load_onnx(model=model_path)
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # 检查模型输入输出信息
    print('--> Model info:')
    print('Inputs:', rknn.inputs)
    print('Outputs:', rknn.outputs)
    print('Input number:', rknn.input_num)
    print('Output number:', rknn.output_num)

    # Build model
    print('--> Building model')
    try:
        ret = rknn.build(
            do_quantization=do_quant, 
            dataset=DATASET_PATH,
            # 添加构建选项以避免段错误
            pre_compile=False,
            # 禁用可能引起问题的优化
            remove_no_pass_ops=False,
            # 使用浮点计算
            float_dtype='float16'
        )
        if ret != 0:
            print('Build model failed! Error code:', ret)
            exit(ret)
        print('done')
    except Exception as e:
        print(f'Build model failed with exception: {e}')
        # 尝试不使用量化
        if do_quant:
            print('Trying without quantization...')
            ret = rknn.build(do_quantization=False, dataset=DATASET_PATH)
            if ret != 0:
                print('Build without quantization also failed!')
                exit(ret)
            print('Build succeeded without quantization')
        else:
            exit(1)

    # Export rknn model
    print('--> Export rknn model')
    ret = rknn.export_rknn(output_path)
    if ret != 0:
        print('Export rknn model failed!')
        exit(ret)
    print('done')

    # Release
    rknn.release()
    print('Model converted successfully!')
